from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import joblib
import pandas as pd
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

def train_and_evaluate_models(X_train, y_train, selected_feature_names):
    """训练并评估多个模型"""
    # 划分训练集和验证集
    X_train, X_val, y_train, y_val = train_test_split(
        X_train, y_train, test_size=0.2, random_state=42, stratify=y_train)

    # 初始化模型
    models = {
        '随机森林': RandomForestClassifier(n_estimators=100, random_state=42, class_weight='balanced'),
        '梯度提升树': GradientBoostingClassifier(n_estimators=100, random_state=42),
        '支持向量机': SVC(probability=True, random_state=42, class_weight='balanced')
    }

    # 训练和评估
    results = {}
    for name, model in models.items():
        print(f"\n正在训练 {name}...")
        model.fit(X_train, y_train)

        # 预测概率
        y_pred_proba = model.predict_proba(X_val)[:, 1]

        # 计算AUC
        auc = roc_auc_score(y_val, y_pred_proba)

        # 存储结果
        results[name] = {
            'model': model,
            'auc': auc
        }

        # 打印结果
        print(f"{name} - AUC分数: {auc:.4f}")

        # 分类报告
        y_pred = model.predict(X_val)
        print(f"\n{name}分类报告:")
        print(classification_report(y_val, y_pred))

        # 混淆矩阵
        cm = confusion_matrix(y_val, y_pred)
        plt.figure(figsize=(6, 4))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
        plt.title(f'混淆矩阵 - {name}')
        plt.xlabel('预测值')
        plt.ylabel('真实值')
        plt.show()

    # 选择最佳模型
    best_model_name = max(results, key=lambda x: results[x]['auc'])
    best_model = results[best_model_name]['model']
    print(f"\n最佳模型: {best_model_name}, AUC分数: {results[best_model_name]['auc']:.4f}")

    # 保存最佳模型
    joblib.dump(best_model, 'best_model.pkl')

    # 特征重要性分析(适用于树模型)
    if hasattr(best_model, 'feature_importances_'):
        feature_importance = pd.DataFrame({
            '特征': selected_feature_names,
            '重要性': best_model.feature_importances_
        }).sort_values('重要性', ascending=False)

        plt.figure(figsize=(10, 8))
        sns.barplot(x='重要性', y='特征', data=feature_importance.head(20))
        plt.title('前20重要特征')
        plt.tight_layout()
        plt.show()

        # 保存特征重要性
        feature_importance.to_csv('feature_importance.csv', index=False)

    return best_model